from pro_data import DataSets
import xgboost as xgb
from utils.utiils_perf import *

def xgb_func(data_train, data_test, label_train):


    dl_train = xgb.DMatrix(data=data_train, label=label_train)
    dl_test = xgb.DMatrix(data=data_test)

    # 参数设置
    params = {'booster': 'gbtree',
              'objective': 'multi:softmax',
              'eval_metric': 'mlogloss',
              'max_depth': 8,
              'lambda': 1,
              'subsample': 0.75,
              'colsample_bytree': 0.75,
              'min_child_weight': 1,
              'eta': 0.025,
              'gamma': 0.4,
              'seed': 0,
              'nthread': 8,
              'silent': True,
              'reg_alpha': 0.1,
              'num_class': 40}

    watchlist = [(dl_train, 'train')]

    # test
    bst = xgb.train(params, dl_train, num_boost_round=500, evals=watchlist)
    pred_test = bst.predict(dl_test)

    return pred_test


if __name__ == "__main__":
    path_x = "../data/data2_naca0012/dv.csv"
    path_y = "../data/data2_naca0012/fc.csv"

    ds = DataSets(path_x=path_x,
                  path_y=path_y,
                  label_index=2,
                  label_threshold=-0.16,
                  test_proportion=0.1)

    loop_cls(ds, xgb_func, 2)

